{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8690ce9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _components.meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f538d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "import builtins\n",
    "import copy as cp\n",
    "import functools\n",
    "import inspect\n",
    "import sys\n",
    "import types\n",
    "from functools import partial, wraps\n",
    "from types import *\n",
    "from typing import *\n",
    "\n",
    "import docstring_parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c59ed2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "\n",
    "from fastkafka._aiokafka_imports import AIOKafkaConsumer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e49f29f9",
   "metadata": {},
   "source": [
    "# Fastcore replacement: patch & delegates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fcc9993",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |exporti\n",
    "\n",
    "\n",
    "def test_eq(a: Any, b: Any) -> None:\n",
    "    \"`test` that `a==b`\"\n",
    "    if a != b:\n",
    "        raise ValueError(f\"{a} != {b}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dd42876",
   "metadata": {},
   "source": [
    "## Patching\n",
    "\n",
    "> copied from https://github.com/fastai/fastcore/blob/master/nbs/01_basics.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a621b3e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |exporti\n",
    "F = TypeVar(\"F\", bound=Callable[..., Any])\n",
    "\n",
    "\n",
    "def copy_func(f: Union[F, FunctionType]) -> Union[F, FunctionType]:\n",
    "    \"Copy a non-builtin function (NB `copy.copy` does not work for this)\"\n",
    "    if not isinstance(f, FunctionType):\n",
    "        return cp.copy(f)\n",
    "    fn = FunctionType(\n",
    "        f.__code__, f.__globals__, f.__name__, f.__defaults__, f.__closure__\n",
    "    )\n",
    "    fn.__kwdefaults__ = f.__kwdefaults__\n",
    "    fn.__dict__.update(f.__dict__)\n",
    "    fn.__annotations__.update(f.__annotations__)\n",
    "    fn.__qualname__ = f.__qualname__\n",
    "    fn.__doc__ = f.__doc__\n",
    "    return fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d465db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def foo():\n",
    "    \"\"\"Test doc\"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "a = cp.copy(foo)\n",
    "b = cp.deepcopy(foo)\n",
    "\n",
    "a.someattr = \"hello\"  # since a and b point at the same object, updating a will update b\n",
    "test_eq(b.someattr, \"hello\")\n",
    "\n",
    "assert a is foo and b is foo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "739b940f",
   "metadata": {},
   "source": [
    "However, with copy_func, you can retrieve a copy of a function without a reference to the original object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db9c3670",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = copy_func(foo)  # c is an indpendent object\n",
    "assert c is not foo\n",
    "assert c.__doc__ == \"\"\"Test doc\"\"\", c.__doc__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88934909",
   "metadata": {},
   "outputs": [],
   "source": [
    "def g(x, *, y=3):\n",
    "    return x + y\n",
    "\n",
    "\n",
    "test_eq(copy_func(g)(4), 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00f7e9ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |exporti\n",
    "\n",
    "\n",
    "def patch_to(\n",
    "    cls: Union[Type, Iterable[Type]], as_prop: bool = False, cls_method: bool = False\n",
    ") -> Callable[[F], F]:\n",
    "    \"Decorator: add `f` to `cls`\"\n",
    "    if not isinstance(cls, (tuple, list)):\n",
    "        cls = (cls,)  # type: ignore\n",
    "\n",
    "    def _inner(f: F) -> F:\n",
    "        for c_ in cls:\n",
    "            nf = copy_func(f)\n",
    "            nm = f.__name__\n",
    "            # `functools.update_wrapper` when passing patched function to `Pipeline`, so we do it manually\n",
    "            for o in functools.WRAPPER_ASSIGNMENTS:\n",
    "                setattr(nf, o, getattr(f, o))\n",
    "            nf.__qualname__ = f\"{c_.__name__}.{nm}\"\n",
    "            if cls_method:\n",
    "                setattr(c_, nm, MethodType(nf, c_))\n",
    "            else:\n",
    "                setattr(c_, nm, property(nf) if as_prop else nf)\n",
    "        # Avoid clobbering existing functions\n",
    "        # nosemgrep\n",
    "        existing_func = globals().get(nm, builtins.__dict__.get(nm, None))\n",
    "        return existing_func  # type: ignore\n",
    "\n",
    "    return _inner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee9ab22b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T3(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "@patch_to(_T3)\n",
    "def foo(self):\n",
    "    \"\"\"Test doc\"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "assert _T3.foo.__doc__ == \"\"\"Test doc\"\"\", foo.__doc__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "740507c5",
   "metadata": {},
   "source": [
    "     \n",
    "The @patch_to decorator allows you to monkey patch a function into a class as a method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce7b854a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T3(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "@patch_to(_T3)\n",
    "def func1(self, a):\n",
    "    return self + a\n",
    "\n",
    "\n",
    "t = _T3(1)  # we initilized `t` to a type int = 1\n",
    "test_eq(t.func1(2), 3)  # we add 2 to `t`, so 2 + 1 = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d73911c9",
   "metadata": {},
   "source": [
    "     \n",
    "You can access instance properties in the usual way via self:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dfcbf2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T4:\n",
    "    def __init__(self, g):\n",
    "        self.g = g\n",
    "\n",
    "\n",
    "@patch_to(_T4)\n",
    "def greet(self, x):\n",
    "    return self.g + x\n",
    "\n",
    "\n",
    "t = _T4(\"hello \")  # this sets self.g = 'helllo '\n",
    "test_eq(\n",
    "    t.greet(\"world\"), \"hello world\"\n",
    ")  # t.greet('world') will append 'world' to 'hello '"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d6df22e",
   "metadata": {},
   "source": [
    "     \n",
    "You can instead specify that the method should be a class method by setting cls_method=True:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3844687b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T5(int):\n",
    "    attr = 3  # attr is a class attribute we will access in a later method\n",
    "\n",
    "\n",
    "@patch_to(_T5, cls_method=True)\n",
    "def func(cls, x):\n",
    "    return cls.attr + x  # you can access class attributes in the normal way\n",
    "\n",
    "\n",
    "test_eq(_T5.func(4), 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4445ccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Additionally you can specify that the function you want to patch should be a class attribute with as_prop=True:\n",
    "\n",
    "\n",
    "@patch_to(_T5, as_prop=True)\n",
    "def add_ten(self):\n",
    "    return self + 10\n",
    "\n",
    "\n",
    "t = _T5(4)\n",
    "test_eq(t.add_ten, 14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a1b91e7",
   "metadata": {},
   "source": [
    "     \n",
    "Instead of passing one class to the @patch_to decorator, you can pass multiple classes in a tuple to simulteanously patch more than one class with the same method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a626f18",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T6(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "class _T7(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "@patch_to((_T6, _T7))\n",
    "def func_mult(self, a):\n",
    "    return self * a\n",
    "\n",
    "\n",
    "t = _T6(2)\n",
    "test_eq(t.func_mult(4), 8)\n",
    "t = _T7(2)\n",
    "test_eq(t.func_mult(4), 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89a579b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exporti\n",
    "\n",
    "\n",
    "def eval_type(\n",
    "    t: Sequence, glb: Optional[Dict[str, Any]], loc: Optional[Mapping[str, object]]\n",
    ") -> Any:\n",
    "    \"`eval` a type or collection of types, if needed, for annotations in py3.10+\"\n",
    "    if isinstance(t, str):\n",
    "        if \"|\" in t:\n",
    "            return Union[eval_type(tuple(t.split(\"|\")), glb, loc)]\n",
    "        # nosemgrep\n",
    "        return eval(t, glb, loc)  # nosec B307:blacklist\n",
    "    if isinstance(t, (tuple, list)):\n",
    "        return type(t)([eval_type(c, glb, loc) for c in t])\n",
    "    return t\n",
    "\n",
    "\n",
    "def union2tuple(t) -> Tuple[Any, ...]:  # type: ignore\n",
    "    if getattr(t, \"__origin__\", None) is Union:\n",
    "        return t.__args__  # type: ignore\n",
    "\n",
    "    if sys.version_info >= (3, 10):\n",
    "        if isinstance(t, UnionType):\n",
    "            return t.__args__\n",
    "\n",
    "    return t  # type: ignore\n",
    "\n",
    "\n",
    "def get_annotations_ex(\n",
    "    obj: Union[FunctionType, Type, F],\n",
    "    *,\n",
    "    globals: Optional[Dict[str, Any]] = None,\n",
    "    locals: Optional[Dict[str, Any]] = None,\n",
    ") -> Tuple[Dict[str, Any], Union[Any, Dict[str, Any], None], Dict[str, Any]]:\n",
    "    \"Backport of py3.10 `get_annotations` that returns globals/locals\"\n",
    "    if isinstance(obj, type):\n",
    "        obj_dict = getattr(obj, \"__dict__\", None)\n",
    "        if obj_dict and hasattr(obj_dict, \"get\"):\n",
    "            ann = obj_dict.get(\"__annotations__\", None)\n",
    "            if isinstance(ann, types.GetSetDescriptorType):\n",
    "                ann = None\n",
    "        else:\n",
    "            ann = None\n",
    "\n",
    "        obj_globals = None\n",
    "        module_name = getattr(obj, \"__module__\", None)\n",
    "        if module_name:\n",
    "            module = sys.modules.get(module_name, None)\n",
    "            if module:\n",
    "                obj_globals = getattr(module, \"__dict__\", None)\n",
    "        obj_locals = dict(vars(obj))\n",
    "        unwrap = obj\n",
    "    elif isinstance(obj, types.ModuleType):\n",
    "        ann = getattr(obj, \"__annotations__\", None)\n",
    "        obj_globals = getattr(obj, \"__dict__\")\n",
    "        obj_locals, unwrap = None, None\n",
    "    elif callable(obj):\n",
    "        ann = getattr(obj, \"__annotations__\", None)\n",
    "        obj_globals = getattr(obj, \"__globals__\", None)\n",
    "        obj_locals, unwrap = None, obj  # type: ignore\n",
    "    else:\n",
    "        raise TypeError(f\"{obj!r} is not a module, class, or callable.\")\n",
    "\n",
    "    if ann is None:\n",
    "        ann = {}\n",
    "    if not isinstance(ann, dict):\n",
    "        raise ValueError(f\"{obj!r}.__annotations__ is neither a dict nor None\")\n",
    "    if not ann:\n",
    "        ann = {}\n",
    "\n",
    "    if unwrap is not None:\n",
    "        while True:\n",
    "            if hasattr(unwrap, \"__wrapped__\"):\n",
    "                unwrap = unwrap.__wrapped__\n",
    "                continue\n",
    "            if isinstance(unwrap, functools.partial):\n",
    "                unwrap = unwrap.func  # type: ignore\n",
    "                continue\n",
    "            break\n",
    "        if hasattr(unwrap, \"__globals__\"):\n",
    "            obj_globals = unwrap.__globals__\n",
    "\n",
    "    if globals is None:\n",
    "        globals = obj_globals\n",
    "    if locals is None:\n",
    "        locals = obj_locals\n",
    "\n",
    "    return dict(ann), globals, locals  # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "615f96d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def patch(  # type: ignore\n",
    "    f: Optional[F] = None, *, as_prop: bool = False, cls_method: bool = False\n",
    "):\n",
    "    \"Decorator: add `f` to the first parameter's class (based on f's type annotations)\"\n",
    "    if f is None:\n",
    "        return partial(patch, as_prop=as_prop, cls_method=cls_method)\n",
    "    ann, glb, loc = get_annotations_ex(f)\n",
    "    cls = union2tuple(\n",
    "        eval_type(ann.pop(\"cls\") if cls_method else next(iter(ann.values())), glb, loc)\n",
    "    )\n",
    "    return patch_to(cls, as_prop=as_prop, cls_method=cls_method)(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a65ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T8(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "@patch\n",
    "def func(self: _T8, a):\n",
    "    \"\"\"Test doc\"\"\"\n",
    "    return self + a\n",
    "\n",
    "\n",
    "assert _T8.func.__doc__ == \"\"\"Test doc\"\"\", func.__doc__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e45c054c",
   "metadata": {},
   "source": [
    "     \n",
    "@patch is an alternative to @patch_to that allows you similarly monkey patch class(es) by using type annotations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d5f6b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T8(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "@patch\n",
    "def func(self: _T8, a):\n",
    "    return self + a\n",
    "\n",
    "\n",
    "t = _T8(1)  # we initilized `t` to a type int = 1\n",
    "test_eq(t.func(3), 4)  # we add 3 to `t`, so 3 + 1 = 4\n",
    "test_eq(t.func.__qualname__, \"_T8.func\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dae3414e",
   "metadata": {},
   "source": [
    "Similarly to patch_to, you can supply a union of classes instead of a single class in your type annotations to patch multiple classes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "367a359d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T9(int):\n",
    "    pass\n",
    "\n",
    "\n",
    "@patch\n",
    "def func2(x: Union[_T8, _T9], a):\n",
    "    return x * a  # will patch both _T8 and _T9\n",
    "\n",
    "\n",
    "t = _T8(2)\n",
    "test_eq(t.func2(4), 8)\n",
    "test_eq(t.func2.__qualname__, \"_T8.func2\")\n",
    "\n",
    "t = _T9(2)\n",
    "test_eq(t.func2(4), 8)\n",
    "test_eq(t.func2.__qualname__, \"_T9.func2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dbb3d5a",
   "metadata": {},
   "source": [
    "     \n",
    "Just like patch_to decorator you can use as_prop and cls_method parameters with patch decorator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eed1b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@patch(as_prop=True)\n",
    "def add_ten(self: _T5):\n",
    "    return self + 10\n",
    "\n",
    "\n",
    "t = _T5(4)\n",
    "test_eq(t.add_ten, 14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a1a6407",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T5(int):\n",
    "    attr = 3  # attr is a class attribute we will access in a later method\n",
    "\n",
    "\n",
    "@patch(cls_method=True)\n",
    "def func(cls: _T5, x):\n",
    "    return cls.attr + x  # you can access class attributes in the normal way\n",
    "\n",
    "\n",
    "test_eq(_T5.func(4), 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "391aec1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95cafec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_sig(f, b):\n",
    "    \"Test the signature of an object\"\n",
    "    if str(inspect.signature(f)) != b:\n",
    "        raise ValueError(f\"{inspect.signature(f)} != {b}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "855fdd0a",
   "metadata": {},
   "source": [
    "# Fastcore meta deps\n",
    "\n",
    "> Copied from https://github.com/fastai/fastcore/blob/master/nbs/07_meta.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb9c8b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |export\n",
    "\n",
    "\n",
    "def _delegates_without_docs(\n",
    "    to: Optional[F] = None,  # Delegatee\n",
    "    keep: bool = False,  # Keep `kwargs` in decorated function?\n",
    "    but: Optional[List[str]] = None,  # Exclude these parameters from signature\n",
    ") -> Callable[[F], F]:\n",
    "    \"Decorator: replace `**kwargs` in signature with params from `to`\"\n",
    "    if but is None:\n",
    "        but = []\n",
    "\n",
    "    def _f(f: F) -> F:\n",
    "        if to is None:\n",
    "            to_f, from_f = f.__base__.__init__, f.__init__  # type: ignore\n",
    "        else:\n",
    "            to_f, from_f = to.__init__ if isinstance(to, type) else to, f  # type: ignore\n",
    "        from_f = getattr(from_f, \"__func__\", from_f)\n",
    "        to_f = getattr(to_f, \"__func__\", to_f)\n",
    "        if hasattr(from_f, \"__delwrap__\"):\n",
    "            return f\n",
    "        sig = inspect.signature(from_f)\n",
    "        sigd = dict(sig.parameters)\n",
    "        if \"kwargs\" in sigd:\n",
    "            k = sigd.pop(\"kwargs\")\n",
    "        else:\n",
    "            k = None\n",
    "        s2 = {\n",
    "            k: v.replace(kind=inspect.Parameter.KEYWORD_ONLY)\n",
    "            for k, v in inspect.signature(to_f).parameters.items()\n",
    "            if v.default != inspect.Parameter.empty and k not in sigd and k not in but  # type: ignore\n",
    "        }\n",
    "        anno = {\n",
    "            k: v\n",
    "            for k, v in getattr(to_f, \"__annotations__\", {}).items()\n",
    "            if k not in sigd and k not in but  # type: ignore\n",
    "        }\n",
    "        sigd.update(s2)\n",
    "        if keep and k is not None:\n",
    "            sigd[\"kwargs\"] = k\n",
    "        else:\n",
    "            from_f.__delwrap__ = to_f\n",
    "        from_f.__signature__ = sig.replace(parameters=list(sigd.values()))\n",
    "        if hasattr(from_f, \"__annotations__\"):\n",
    "            from_f.__annotations__.update(anno)\n",
    "        return f\n",
    "\n",
    "    return _f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cb1c2e5",
   "metadata": {},
   "source": [
    "A common Python idiom is to accept **kwargs in addition to named parameters that are passed onto other function calls. It is especially common to use **kwargs when you want to give the user an option to override default parameters of any functions or methods being called by the parent function.\n",
    "\n",
    "For example, suppose we have have a function foo that passes arguments to baz like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a80423a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def baz(a, b: int = 2, c: int = 3) -> int:\n",
    "    \"\"\"Baz\n",
    "    Params:\n",
    "        a: something\n",
    "        b: whatever\n",
    "        c: whocares\n",
    "\n",
    "    Returns:\n",
    "        Nada\n",
    "    \"\"\"\n",
    "    return a + b + c\n",
    "\n",
    "\n",
    "def foo(c, a, **kwargs):\n",
    "    return c + baz(a, **kwargs)\n",
    "\n",
    "\n",
    "assert foo(c=1, a=1) == 7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9acda69",
   "metadata": {},
   "source": [
    "The problem with this approach is the api for foo is obfuscated. Users cannot introspect what the valid arguments for **kwargs are without reading the source code. When a user tries tries to introspect the signature of foo, they are presented with this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad04fa0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (c, a, **kwargs)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inspect.signature(foo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1a2627f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (a, b: int = 2, c: int = 3) -> int>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inspect.signature(baz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c167e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Baz\n",
      "    Params:\n",
      "        a: something\n",
      "        b: whatever\n",
      "        c: whocares\n",
      "\n",
      "    Returns:\n",
      "        Nada\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "print(baz.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1668d82",
   "metadata": {},
   "source": [
    "We can address this issue by using the decorator delegates to include parameters from other functions. For example, if we apply the delegates decorator to foo to include parameters from baz:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de7978e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (c, a, *, b: int = 2)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@_delegates_without_docs(baz)\n",
    "def foo(c, a, **kwargs):\n",
    "    \"\"\"Foo is great\n",
    "\n",
    "    Params:\n",
    "        c: c from foo\n",
    "        a: a from foo\n",
    "    \"\"\"\n",
    "    return c + baz(a, **kwargs)\n",
    "\n",
    "\n",
    "test_sig(foo, \"(c, a, *, b: int = 2)\")\n",
    "assert (\n",
    "    foo.__doc__\n",
    "    == \"Foo is great\\n\\n    Params:\\n        c: c from foo\\n        a: a from foo\\n    \"\n",
    ")\n",
    "inspect.signature(foo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd13be32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Foo is great\\n\\n    Params:\\n        c: c from foo\\n        a: a from foo\\n    '"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foo.__doc__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d4b88e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _format_args(xs: List[docstring_parser.DocstringParam]) -> str:\n",
    "    return \"\\nArgs:\\n - \" + \"\\n - \".join(\n",
    "        [f\"{x.arg_name} ({x.type_name}): {x.description}\" for x in xs]\n",
    "    )\n",
    "\n",
    "\n",
    "def combine_params(\n",
    "    f: F, o: Union[Type, Callable[..., Any]], but: Optional[List[str]] = None\n",
    ") -> F:\n",
    "    \"\"\"Combines docstring arguments of a function and another object or function\n",
    "\n",
    "    Args:\n",
    "        f: destination functions where combined arguments will end up\n",
    "        o: source function from which arguments are taken from\n",
    "\n",
    "    Returns:\n",
    "        Function f with augumented docstring including arguments from both functions/objects\n",
    "    \"\"\"\n",
    "    if but is None:\n",
    "        but = []\n",
    "\n",
    "    src_params = docstring_parser.parse_from_object(o).params\n",
    "    #     logger.info(f\"combine_params(): source:{_format_args(src_params)}\")\n",
    "    docs = docstring_parser.parse_from_object(f)\n",
    "    #     logger.info(f\"combine_params(): destination:{_format_args(docs.params)}\")\n",
    "    dst_params_names = [p.arg_name for p in docs.params]\n",
    "\n",
    "    combined_params = docs.params + [\n",
    "        x\n",
    "        for x in src_params\n",
    "        if x.arg_name not in dst_params_names and x.arg_name not in but\n",
    "    ]\n",
    "    #     logger.info(f\"combine_params(): combined:{_format_args(combined_params)}\")\n",
    "\n",
    "    docs.meta = [\n",
    "        x for x in docs.meta if not isinstance(x, docstring_parser.DocstringParam)\n",
    "    ] + combined_params  # type: ignore\n",
    "\n",
    "    f.__doc__ = docstring_parser.compose(\n",
    "        docs, style=docstring_parser.DocstringStyle.GOOGLE\n",
    "    )\n",
    "    return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c76e597d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f2(a: int = 0, b: str = \"nada\"):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        a: parameter a\n",
    "        b: parameter bbbb\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "def f1(b: str, c: int, **kwargs):\n",
    "    \"\"\"Function f1\n",
    "    Args:\n",
    "        b: parameter b\n",
    "        c: parameter c\n",
    "\n",
    "    Raises:\n",
    "        ValueError: sometimes\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "combine_params(f1, f2).__doc__\n",
    "\n",
    "expected = \"\"\"Function f1\n",
    "Args:\n",
    "    b: parameter b\n",
    "    c: parameter c\n",
    "    a: parameter a\n",
    "\n",
    "Raises:\n",
    "    ValueError: sometimes\"\"\"\n",
    "\n",
    "assert f1.__doc__ == expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e98db95a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add test case to test combine_params with 'but' param\n",
    "\n",
    "\n",
    "def f2(a: int = 0, b: str = \"nada\", d: str = \"dada\"):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        a: parameter a\n",
    "        b: parameter bbbb\n",
    "        d: parameter d\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "def f1(b: str, c: int):\n",
    "    \"\"\"Function f1\n",
    "    Args:\n",
    "        b: parameter b\n",
    "        c: parameter c\n",
    "\n",
    "    Raises:\n",
    "        ValueError: sometimes\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "combine_params(f1, f2, but=[\"d\"]).__doc__\n",
    "\n",
    "expected = \"\"\"Function f1\n",
    "Args:\n",
    "    b: parameter b\n",
    "    c: parameter c\n",
    "    a: parameter a\n",
    "\n",
    "Raises:\n",
    "    ValueError: sometimes\"\"\"\n",
    "\n",
    "assert f1.__doc__ == expected, f1.__doc__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a06f975f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def delegates(\n",
    "    o: Union[Type, Callable[..., Any]],\n",
    "    keep: bool = False,\n",
    "    but: Optional[List[str]] = None,\n",
    ") -> Callable[[F], F]:\n",
    "    \"\"\"Delegates keyword agruments from o to the function the decorator is applied to\n",
    "\n",
    "    Args:\n",
    "        o: object (class or function) with default kwargs\n",
    "        keep: Keep `kwargs` in decorated function?\n",
    "        but: argument names not to include\n",
    "    \"\"\"\n",
    "\n",
    "    def _inner(f: F, keep: bool = keep, but: Optional[List[str]] = but) -> F:\n",
    "        def _combine_params(\n",
    "            o: Union[Type, Callable[..., Any]], but: Optional[List[str]] = None\n",
    "        ) -> Callable[[F], F]:\n",
    "            def __combine_params(\n",
    "                f: F,\n",
    "                o: Union[Type, Callable[..., Any]] = o,\n",
    "                but: Optional[List[str]] = but,\n",
    "            ) -> F:\n",
    "                return combine_params(f=f, o=o, but=but)\n",
    "\n",
    "            return __combine_params\n",
    "\n",
    "        @_combine_params(o, but=but)  # type: ignore\n",
    "        @_delegates_without_docs(o, keep=keep, but=but)  # type: ignore\n",
    "        @wraps(f)\n",
    "        def _f(*args: Any, **kwargs: Any) -> Any:\n",
    "            return f(*args, **kwargs)\n",
    "\n",
    "        return _f\n",
    "\n",
    "    return _inner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "827294c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f2(a: str = \"whatever\", d: int = 42) -> None:\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        a: parameter a\n",
    "        b: parameter bbbb\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "@delegates(f2)\n",
    "def f1(b: str, c: int, **kwargs):\n",
    "    \"\"\"Function f1\n",
    "    Args:\n",
    "        b: parameter b\n",
    "        c: parameter c\n",
    "\n",
    "    Raises:\n",
    "        ValueError: sometimes\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "expected = \"\"\"Function f1\n",
    "Args:\n",
    "    b: parameter b\n",
    "    c: parameter c\n",
    "    a: parameter a\n",
    "\n",
    "Raises:\n",
    "    ValueError: sometimes\"\"\"\n",
    "\n",
    "assert f1.__doc__ == expected\n",
    "assert len(inspect.signature(f2).parameters) == 2\n",
    "assert len(inspect.signature(f1).parameters) == 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4443b59d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add test case to test delegates with 'but' param\n",
    "\n",
    "\n",
    "def p1(a: int = 0, b: float = 0.1):\n",
    "    \"\"\"Func p1\n",
    "    Args:\n",
    "        a: hello\n",
    "        b: bello\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "@delegates(p1, but=[\"b\"])\n",
    "def p2(c: int, d: float):\n",
    "    \"\"\"Func p2\n",
    "    Args:\n",
    "        c: cello\n",
    "        d: dello\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "expected = \"\"\"Func p2\n",
    "Args:\n",
    "    c: cello\n",
    "    d: dello\n",
    "    a: hello\"\"\"\n",
    "\n",
    "assert p2.__doc__ == expected, p2.__doc__\n",
    "assert len(inspect.signature(p2).parameters) == 3\n",
    "assert len(inspect.signature(p1).parameters) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a06d70",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mappingproxy({'b': <Parameter \"b: str\">,\n",
       "              'c': <Parameter \"c: int\">,\n",
       "              'a': <Parameter \"a: str = 'whatever'\">,\n",
       "              'd': <Parameter \"d: int = 42\">})"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@delegates(f2)\n",
    "def f3(b: str, c: int, **kwargs):\n",
    "    \"\"\"Function f1\n",
    "    Args:\n",
    "        b: parameter b\n",
    "        c: parameter c\n",
    "\n",
    "    Raises:\n",
    "        ValueError: sometimes\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "params = inspect.signature(f3).parameters\n",
    "display(params)\n",
    "assert len(params) == 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ea7197",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "function a\n",
      "\n",
      "Args:\n",
      "    a: parameter a\n",
      "    *topics (list(str)): optional list of topics to subscribe to. If not set,\n",
      "        call :meth:`.subscribe` or :meth:`.assign` before consuming records.\n",
      "        Passing topics directly is same as calling :meth:`.subscribe` API.\n",
      "    bootstrap_servers (str, list(str)): a ``host[:port]`` string (or list of\n",
      "        ``host[:port]`` strings) that the consumer should contact to bootstrap\n",
      "        initial cluster metadata.\n",
      "        \n",
      "        This does not have to be the full node list.\n",
      "        It just needs to have at least one broker that will respond to a\n",
      "        Metadata API Request. Default port is 9092. If no servers are\n",
      "        specified, will default to ``localhost:9092``.\n",
      "    client_id (str): a name for this client. This string is passed in\n",
      "        each request to servers and can be used to identify specific\n",
      "        server-side log entries that correspond to this client. Also\n",
      "        submitted to :class:`~.consumer.group_coordinator.GroupCoordinator`\n",
      "        for logging with respect to consumer group administration. Default:\n",
      "        ``aiokafka-{version}``\n",
      "    group_id (str or None): name of the consumer group to join for dynamic\n",
      "        partition assignment (if enabled), and to use for fetching and\n",
      "        committing offsets. If None, auto-partition assignment (via\n",
      "        group coordinator) and offset commits are disabled.\n",
      "        Default: None\n",
      "    key_deserializer (Callable): Any callable that takes a\n",
      "        raw message key and returns a deserialized key.\n",
      "    value_deserializer (Callable, Optional): Any callable that takes a\n",
      "        raw message value and returns a deserialized value.\n",
      "    fetch_min_bytes (int): Minimum amount of data the server should\n",
      "        return for a fetch request, otherwise wait up to\n",
      "        `fetch_max_wait_ms` for more data to accumulate. Default: 1.\n",
      "    fetch_max_bytes (int): The maximum amount of data the server should\n",
      "        return for a fetch request. This is not an absolute maximum, if\n",
      "        the first message in the first non-empty partition of the fetch\n",
      "        is larger than this value, the message will still be returned\n",
      "        to ensure that the consumer can make progress. NOTE: consumer\n",
      "        performs fetches to multiple brokers in parallel so memory\n",
      "        usage will depend on the number of brokers containing\n",
      "        partitions for the topic.\n",
      "        Supported Kafka version >= 0.10.1.0. Default: 52428800 (50 Mb).\n",
      "    fetch_max_wait_ms (int): The maximum amount of time in milliseconds\n",
      "        the server will block before answering the fetch request if\n",
      "        there isn't sufficient data to immediately satisfy the\n",
      "        requirement given by fetch_min_bytes. Default: 500.\n",
      "    max_partition_fetch_bytes (int): The maximum amount of data\n",
      "        per-partition the server will return. The maximum total memory\n",
      "        used for a request ``= #partitions * max_partition_fetch_bytes``.\n",
      "        This size must be at least as large as the maximum message size\n",
      "        the server allows or else it is possible for the producer to\n",
      "        send messages larger than the consumer can fetch. If that\n",
      "        happens, the consumer can get stuck trying to fetch a large\n",
      "        message on a certain partition. Default: 1048576.\n",
      "    max_poll_records (int): The maximum number of records returned in a\n",
      "        single call to :meth:`.getmany`. Defaults ``None``, no limit.\n",
      "    request_timeout_ms (int): Client request timeout in milliseconds.\n",
      "        Default: 40000.\n",
      "    retry_backoff_ms (int): Milliseconds to backoff when retrying on\n",
      "        errors. Default: 100.\n",
      "    auto_offset_reset (str): A policy for resetting offsets on\n",
      "        :exc:`.OffsetOutOfRangeError` errors: ``earliest`` will move to the oldest\n",
      "        available message, ``latest`` will move to the most recent, and\n",
      "        ``none`` will raise an exception so you can handle this case.\n",
      "        Default: ``latest``.\n",
      "    enable_auto_commit (bool): If true the consumer's offset will be\n",
      "        periodically committed in the background. Default: True.\n",
      "    auto_commit_interval_ms (int): milliseconds between automatic\n",
      "        offset commits, if enable_auto_commit is True. Default: 5000.\n",
      "    check_crcs (bool): Automatically check the CRC32 of the records\n",
      "        consumed. This ensures no on-the-wire or on-disk corruption to\n",
      "        the messages occurred. This check adds some overhead, so it may\n",
      "        be disabled in cases seeking extreme performance. Default: True\n",
      "    metadata_max_age_ms (int): The period of time in milliseconds after\n",
      "        which we force a refresh of metadata even if we haven't seen any\n",
      "        partition leadership changes to proactively discover any new\n",
      "        brokers or partitions. Default: 300000\n",
      "    partition_assignment_strategy (list): List of objects to use to\n",
      "        distribute partition ownership amongst consumer instances when\n",
      "        group management is used. This preference is implicit in the order\n",
      "        of the strategies in the list. When assignment strategy changes:\n",
      "        to support a change to the assignment strategy, new versions must\n",
      "        enable support both for the old assignment strategy and the new\n",
      "        one. The coordinator will choose the old assignment strategy until\n",
      "        all members have been updated. Then it will choose the new\n",
      "        strategy. Default: [:class:`.RoundRobinPartitionAssignor`]\n",
      "    max_poll_interval_ms (int): Maximum allowed time between calls to\n",
      "        consume messages (e.g., :meth:`.getmany`). If this interval\n",
      "        is exceeded the consumer is considered failed and the group will\n",
      "        rebalance in order to reassign the partitions to another consumer\n",
      "        group member. If API methods block waiting for messages, that time\n",
      "        does not count against this timeout. See `KIP-62`_ for more\n",
      "        information. Default 300000\n",
      "    rebalance_timeout_ms (int): The maximum time server will wait for this\n",
      "        consumer to rejoin the group in a case of rebalance. In Java client\n",
      "        this behaviour is bound to `max.poll.interval.ms` configuration,\n",
      "        but as ``aiokafka`` will rejoin the group in the background, we\n",
      "        decouple this setting to allow finer tuning by users that use\n",
      "        :class:`.ConsumerRebalanceListener` to delay rebalacing. Defaults\n",
      "        to ``session_timeout_ms``\n",
      "    session_timeout_ms (int): Client group session and failure detection\n",
      "        timeout. The consumer sends periodic heartbeats\n",
      "        (`heartbeat.interval.ms`) to indicate its liveness to the broker.\n",
      "        If no hearts are received by the broker for a group member within\n",
      "        the session timeout, the broker will remove the consumer from the\n",
      "        group and trigger a rebalance. The allowed range is configured with\n",
      "        the **broker** configuration properties\n",
      "        `group.min.session.timeout.ms` and `group.max.session.timeout.ms`.\n",
      "        Default: 10000\n",
      "    heartbeat_interval_ms (int): The expected time in milliseconds\n",
      "        between heartbeats to the consumer coordinator when using\n",
      "        Kafka's group management feature. Heartbeats are used to ensure\n",
      "        that the consumer's session stays active and to facilitate\n",
      "        rebalancing when new consumers join or leave the group. The\n",
      "        value must be set lower than `session_timeout_ms`, but typically\n",
      "        should be set no higher than 1/3 of that value. It can be\n",
      "        adjusted even lower to control the expected time for normal\n",
      "        rebalances. Default: 3000\n",
      "    consumer_timeout_ms (int): maximum wait timeout for background fetching\n",
      "        routine. Mostly defines how fast the system will see rebalance and\n",
      "        request new data for new partitions. Default: 200\n",
      "    api_version (str): specify which kafka API version to use.\n",
      "        :class:`AIOKafkaConsumer` supports Kafka API versions >=0.9 only.\n",
      "        If set to ``auto``, will attempt to infer the broker version by\n",
      "        probing various APIs. Default: ``auto``\n",
      "    security_protocol (str): Protocol used to communicate with brokers.\n",
      "        Valid values are: ``PLAINTEXT``, ``SSL``. Default: ``PLAINTEXT``.\n",
      "    ssl_context (ssl.SSLContext): pre-configured :class:`~ssl.SSLContext`\n",
      "        for wrapping socket connections. Directly passed into asyncio's\n",
      "        :meth:`~asyncio.loop.create_connection`. For more information see\n",
      "        :ref:`ssl_auth`. Default: None.\n",
      "    exclude_internal_topics (bool): Whether records from internal topics\n",
      "        (such as offsets) should be exposed to the consumer. If set to True\n",
      "        the only way to receive records from an internal topic is\n",
      "        subscribing to it. Requires 0.10+ Default: True\n",
      "    connections_max_idle_ms (int): Close idle connections after the number\n",
      "        of milliseconds specified by this config. Specifying `None` will\n",
      "        disable idle checks. Default: 540000 (9 minutes).\n",
      "    isolation_level (str): Controls how to read messages written\n",
      "        transactionally.\n",
      "        \n",
      "        If set to ``read_committed``, :meth:`.getmany` will only return\n",
      "        transactional messages which have been committed.\n",
      "        If set to ``read_uncommitted`` (the default), :meth:`.getmany` will\n",
      "        return all messages, even transactional messages which have been\n",
      "        aborted.\n",
      "        \n",
      "        Non-transactional messages will be returned unconditionally in\n",
      "        either mode.\n",
      "        \n",
      "        Messages will always be returned in offset order. Hence, in\n",
      "        `read_committed` mode, :meth:`.getmany` will only return\n",
      "        messages up to the last stable offset (LSO), which is the one less\n",
      "        than the offset of the first open transaction. In particular any\n",
      "        messages appearing after messages belonging to ongoing transactions\n",
      "        will be withheld until the relevant transaction has been completed.\n",
      "        As a result, `read_committed` consumers will not be able to read up\n",
      "        to the high watermark when there are in flight transactions.\n",
      "        Further, when in `read_committed` the seek_to_end method will\n",
      "        return the LSO. See method docs below. Default: ``read_uncommitted``\n",
      "    sasl_mechanism (str): Authentication mechanism when security_protocol\n",
      "        is configured for ``SASL_PLAINTEXT`` or ``SASL_SSL``. Valid values are:\n",
      "        ``PLAIN``, ``GSSAPI``, ``SCRAM-SHA-256``, ``SCRAM-SHA-512``,\n",
      "        ``OAUTHBEARER``.\n",
      "        Default: ``PLAIN``\n",
      "    sasl_plain_username (str): username for SASL ``PLAIN`` authentication.\n",
      "        Default: None\n",
      "    sasl_plain_password (str): password for SASL ``PLAIN`` authentication.\n",
      "        Default: None\n",
      "    sasl_oauth_token_provider (~aiokafka.abc.AbstractTokenProvider): OAuthBearer token provider instance. (See :mod:`kafka.oauth.abstract`).\n",
      "        Default: None\n",
      "\n",
      "Returns:\n",
      "    : things not stuff\n"
     ]
    }
   ],
   "source": [
    "@delegates(AIOKafkaConsumer)\n",
    "def f(a: int, **kwargs) -> str:\n",
    "    \"\"\"function a\n",
    "\n",
    "    Args:\n",
    "        a: parameter a\n",
    "\n",
    "    Returns:\n",
    "        things not stuff\n",
    "    \"\"\"\n",
    "    print(f\"{a=}\")\n",
    "\n",
    "\n",
    "assert len(f.__doc__) > 5000\n",
    "\n",
    "print(f.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6b11e13",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mappingproxy({'a': <Parameter \"a: int\">,\n",
       "              'loop': <Parameter \"loop=None\">,\n",
       "              'bootstrap_servers': <Parameter \"bootstrap_servers='localhost'\">,\n",
       "              'client_id': <Parameter \"client_id='aiokafka-0.8.0'\">,\n",
       "              'group_id': <Parameter \"group_id=None\">,\n",
       "              'key_deserializer': <Parameter \"key_deserializer=None\">,\n",
       "              'value_deserializer': <Parameter \"value_deserializer=None\">,\n",
       "              'fetch_max_wait_ms': <Parameter \"fetch_max_wait_ms=500\">,\n",
       "              'fetch_max_bytes': <Parameter \"fetch_max_bytes=52428800\">,\n",
       "              'fetch_min_bytes': <Parameter \"fetch_min_bytes=1\">,\n",
       "              'max_partition_fetch_bytes': <Parameter \"max_partition_fetch_bytes=1048576\">,\n",
       "              'request_timeout_ms': <Parameter \"request_timeout_ms=40000\">,\n",
       "              'retry_backoff_ms': <Parameter \"retry_backoff_ms=100\">,\n",
       "              'auto_offset_reset': <Parameter \"auto_offset_reset='latest'\">,\n",
       "              'enable_auto_commit': <Parameter \"enable_auto_commit=True\">,\n",
       "              'auto_commit_interval_ms': <Parameter \"auto_commit_interval_ms=5000\">,\n",
       "              'check_crcs': <Parameter \"check_crcs=True\">,\n",
       "              'metadata_max_age_ms': <Parameter \"metadata_max_age_ms=300000\">,\n",
       "              'partition_assignment_strategy': <Parameter \"partition_assignment_strategy=(<class 'kafka.coordinator.assignors.roundrobin.RoundRobinPartitionAssignor'>,)\">,\n",
       "              'max_poll_interval_ms': <Parameter \"max_poll_interval_ms=300000\">,\n",
       "              'rebalance_timeout_ms': <Parameter \"rebalance_timeout_ms=None\">,\n",
       "              'session_timeout_ms': <Parameter \"session_timeout_ms=10000\">,\n",
       "              'heartbeat_interval_ms': <Parameter \"heartbeat_interval_ms=3000\">,\n",
       "              'consumer_timeout_ms': <Parameter \"consumer_timeout_ms=200\">,\n",
       "              'max_poll_records': <Parameter \"max_poll_records=None\">,\n",
       "              'ssl_context': <Parameter \"ssl_context=None\">,\n",
       "              'security_protocol': <Parameter \"security_protocol='PLAINTEXT'\">,\n",
       "              'api_version': <Parameter \"api_version='auto'\">,\n",
       "              'exclude_internal_topics': <Parameter \"exclude_internal_topics=True\">,\n",
       "              'connections_max_idle_ms': <Parameter \"connections_max_idle_ms=540000\">,\n",
       "              'isolation_level': <Parameter \"isolation_level='read_uncommitted'\">,\n",
       "              'sasl_mechanism': <Parameter \"sasl_mechanism='PLAIN'\">,\n",
       "              'sasl_plain_password': <Parameter \"sasl_plain_password=None\">,\n",
       "              'sasl_plain_username': <Parameter \"sasl_plain_username=None\">,\n",
       "              'sasl_kerberos_service_name': <Parameter \"sasl_kerberos_service_name='kafka'\">,\n",
       "              'sasl_kerberos_domain_name': <Parameter \"sasl_kerberos_domain_name=None\">,\n",
       "              'sasl_oauth_token_provider': <Parameter \"sasl_oauth_token_provider=None\">})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inspect.signature(f).parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5f22fcf",
   "metadata": {},
   "source": [
    "We can optionally decide to keep **kwargs by setting keep=True:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c2fef6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (c, a, *, b: int = 2, **kwargs)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@delegates(baz, keep=True)\n",
    "def foo(c, a, **kwargs):\n",
    "    return c + baz(a, **kwargs)\n",
    "\n",
    "\n",
    "inspect.signature(foo)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53cde9c4",
   "metadata": {},
   "source": [
    "It is important to note that only parameters with default parameters are included. For example, in the below scenario only c, but NOT e and d are included in the signature of foo after applying delegates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ac19ca4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (a, b=1, *, c=2)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def basefoo(e, d, c=2):\n",
    "    pass\n",
    "\n",
    "\n",
    "@delegates(basefoo)\n",
    "def foo(a, b=1, **kwargs):\n",
    "    pass\n",
    "\n",
    "\n",
    "inspect.signature(\n",
    "    foo\n",
    ")  # e and d are not included b/c they don't have default parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b99291",
   "metadata": {},
   "source": [
    "The reason that required arguments (i.e. those without default parameters) are automatically excluded is that you should be explicitly implementing required arguments into your function's signature rather than relying on delegates.\n",
    "\n",
    "Additionally, you can exclude specific parameters from being included in the signature with the but parameter. In the example below, we exclude the parameter d:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90b04168",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (a, b=1, *, c=2)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def basefoo(e, c=2, d=3):\n",
    "    pass\n",
    "\n",
    "\n",
    "@delegates(basefoo, but=[\"d\"])\n",
    "def foo(a, b=1, **kwargs):\n",
    "    pass\n",
    "\n",
    "\n",
    "test_sig(foo, \"(a, b=1, *, c=2)\")\n",
    "inspect.signature(foo)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2961cb40",
   "metadata": {},
   "source": [
    "You can also use delegates between methods in a class. Here is an example of delegates with class methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a23bc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example 1: class methods\n",
    "class _T:\n",
    "    @classmethod\n",
    "    def foo(cls, a=1, b=2):\n",
    "        pass\n",
    "\n",
    "    @classmethod\n",
    "    @delegates(foo)\n",
    "    def bar(cls, c=3, **kwargs):\n",
    "        pass\n",
    "\n",
    "\n",
    "test_sig(_T.bar, \"(c=3, *, a=1, b=2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d2bb41e",
   "metadata": {},
   "source": [
    "Here is the same example with instance methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "789a9cc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example 2: instance methods\n",
    "class _T:\n",
    "    def foo(self, a=1, b=2):\n",
    "        pass\n",
    "\n",
    "    @delegates(foo)\n",
    "    def bar(self, c=3, **kwargs):\n",
    "        pass\n",
    "\n",
    "\n",
    "t = _T()\n",
    "test_sig(t.bar, \"(c=3, *, a=1, b=2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97b9eb07",
   "metadata": {},
   "source": [
    "You can also delegate between classes. By default, the delegates decorator will delegate to the superclass:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d397cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class BaseFoo:\n",
    "#     def __init__(self, e, c=2): pass\n",
    "\n",
    "# @delegates()# since no argument was passsed here we delegate to the superclass\n",
    "# class Foo(BaseFoo):\n",
    "#     def __init__(self, a, b=1, **kwargs): super().__init__(**kwargs)\n",
    "\n",
    "# test_sig(Foo, '(a, b=1, *, c=2)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1d1b68f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def use_parameters_of(\n",
    "    o: Union[Type, Callable[..., Any]], **kwargs: Dict[str, Any]\n",
    ") -> Dict[str, Any]:\n",
    "    \"\"\"Restrict parameters passwed as keyword arguments to parameters from the signature of ``o``\n",
    "\n",
    "    Args:\n",
    "        o: object or callable which signature is used for restricting keyword arguments\n",
    "        kwargs: keyword arguments\n",
    "\n",
    "    Returns:\n",
    "        restricted keyword arguments\n",
    "\n",
    "    \"\"\"\n",
    "    allowed_keys = set(inspect.signature(o).parameters.keys())\n",
    "    return {k: v for k, v in kwargs.items() if k in allowed_keys}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f329c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert use_parameters_of(AIOKafkaConsumer, api_version=0.1, radnom_param=\"random\") == {\n",
    "    \"api_version\": 0.1\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f1013a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def filter_using_signature(f: Callable, **kwargs: Dict[str, Any]) -> Dict[str, Any]:\n",
    "    \"\"\"todo: write docs\"\"\"\n",
    "    param_names = list(inspect.signature(f).parameters.keys())\n",
    "    return {k: v for k, v in kwargs.items() if k in param_names}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9051465c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(a: int, *, b: str):\n",
    "    pass\n",
    "\n",
    "\n",
    "assert filter_using_signature(f, a=1, c=3) == {\"a\": 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1123f286",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "TorF = TypeVar(\"TorF\", Type, Callable[..., Any])\n",
    "\n",
    "\n",
    "def export(module_name: str) -> Callable[[TorF], TorF]:\n",
    "    \"\"\"\n",
    "    Decorator that sets the __module__ attribute of the decorated object to the specified module name.\n",
    "\n",
    "    Args:\n",
    "        module_name: Name of the module to set as __module__ attribute.\n",
    "\n",
    "    Returns:\n",
    "        Decorator function that sets the __module__ attribute of the decorated object.\n",
    "    \"\"\"\n",
    "    def _inner(o: TorF, module_name: str = module_name) -> TorF:\n",
    "        o.__module__ = module_name\n",
    "        return o\n",
    "\n",
    "    return _inner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a09e2c63",
   "metadata": {},
   "outputs": [],
   "source": [
    "@export(\"super.cool\")\n",
    "class A:\n",
    "    pass\n",
    "\n",
    "\n",
    "assert A.__module__ == \"super.cool\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7ca0ffe",
   "metadata": {},
   "source": [
    "# Class context manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce80a4fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "T = TypeVar(\"T\")\n",
    "\n",
    "\n",
    "def classcontextmanager(name: str = \"lifecycle\") -> Callable[[Type[T]], Type[T]]:\n",
    "    \"\"\"\n",
    "    Decorator that adds context manager functionality to a class.\n",
    "\n",
    "    Args:\n",
    "        name: Name of the context manager attribute in the class. Default is \"lifecycle\".\n",
    "\n",
    "    Returns:\n",
    "        Decorator function that adds context manager functionality to the class.\n",
    "    \"\"\"\n",
    "    def _classcontextmanager(cls: Type[T], name: str = name) -> Type[T]:\n",
    "        if not hasattr(cls, name):\n",
    "            raise ValueError\n",
    "\n",
    "        @patch\n",
    "        def __enter__(self: cls) -> Any:  # type: ignore\n",
    "            if not hasattr(self, \"_lifecycle_ctx\"):\n",
    "                self._lifecycle_ctx = []  # type: ignore\n",
    "\n",
    "            self._lifecycle_ctx.append(getattr(self, name)())  # type: ignore\n",
    "            return self._lifecycle_ctx[-1].__enter__()  # type: ignore\n",
    "\n",
    "        @patch\n",
    "        def __exit__(self: cls, *args: Any) -> None:  # type: ignore\n",
    "            self._lifecycle_ctx.pop(-1).__exit__(*args)  # type: ignore\n",
    "\n",
    "        return cls\n",
    "\n",
    "    return _classcontextmanager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e61bdb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@classcontextmanager(\"lifecycle\")\n",
    "class A:\n",
    "    @contextmanager\n",
    "    def lifecycle(self):\n",
    "        try:\n",
    "            print(\"I'm in\")\n",
    "            yield\n",
    "        finally:\n",
    "            print(\"I'm out\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42dbd707",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm in\n",
      "I'm in\n",
      "I'm out\n",
      "I'm out\n"
     ]
    }
   ],
   "source": [
    "a = A()\n",
    "with a:\n",
    "    with a:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80105190",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_default_kwargs_from_sig(f: F, **kwargs: Any) -> Dict[str, Any]:\n",
    "    \"\"\"\n",
    "    Get default values for function **kwargs\n",
    "\n",
    "    Args:\n",
    "        f: Function to extract default values from\n",
    "\n",
    "    Returns:\n",
    "        Dict of default values of function f **kwargs\n",
    "    \"\"\"\n",
    "    defaults = {\n",
    "        k: v.default\n",
    "        for k, v in inspect.signature(f).parameters.items()\n",
    "        if v.default != inspect._empty\n",
    "    }\n",
    "    defaults.update(kwargs)\n",
    "    return defaults"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9dbfd25",
   "metadata": {},
   "outputs": [],
   "source": [
    "def A(c: int, b: str = \"default\", d: float = 3.14):\n",
    "    pass\n",
    "\n",
    "\n",
    "assert _get_default_kwargs_from_sig(A) == {\"b\": \"default\", \"d\": 3.14}\n",
    "assert _get_default_kwargs_from_sig(A, d=5.15) == {\"b\": \"default\", \"d\": 5.15}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54081356",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
